package geppetto.phraseTable.prunning.global;

import java.util.LinkedList;

import geppetto.decoder.TranslationHypothesis;
import geppetto.decoder.TranslationOption;
import geppetto.decoder.scores.GraphPhraseProductScore;
import geppetto.decoder.scores.NumberOfOptionsScore;
import geppetto.phraseHMM.phraseExtraction.extractedphrase.ExtractedPhrasePair;
import geppetto.phraseTable.PhraseTable;
import geppetto.phraseTable.phrase.Phrase;
import geppetto.phraseTable.phrase.feature.LexicalWeightingFeature;
import geppetto.phraseTable.phrase.feature.PenaltyFeature;
import geppetto.phraseTable.phrase.feature.ProbabilityFeature;
import geppetto.phraseTable.phrase.feature.ReorderingFeature;
import geppetto.phraseTable.resource.TranslationGraphResource;
import geppetto.reordering.MSDOrientationCodes;
// this class needs to be generalized somehow 
public class LossBasedPrunner extends AbstractGlobalPhrasePrunner{
	protected final static String PROB_HEADER = "prob";
	protected final static String RPROB_HEADER = "rprob";
	protected final static String LEXICAL_WEIGHITING_HEADER = "lweight";
	protected final static String REVERSE_LEXICAL_WEIGHTING_HEADER = "rlweight";
	protected final static String PENALTY_HEADER = "pen";
	protected final static String REORDERING_HEADER = "reord";

	protected TranslationGraphResource translationGraph;
	protected double threshold;
	protected double[] translationWeights;
	protected double[] reorderingWeights;
	
	public LossBasedPrunner(TranslationGraphResource translationGraph, double[] translationWeights, double[] reorderingWeights , double threshold) {
		super();
		this.translationGraph = translationGraph;
		this.threshold = threshold;
		this.reorderingWeights = reorderingWeights;
		this.translationWeights = translationWeights;
	}

	@Override
	public boolean shouldPrunePhrase(Phrase p, PhraseTable table) {
		LinkedList<TranslationHypothesis> hyps = translationGraph.getFinalHyps();
		NumberOfOptionsScore lengthScore = (NumberOfOptionsScore)translationGraph.getScore(TranslationGraphResource.NUMBER_OF_OPTIONS);
		
		double bestHypsScore = 0;
		TranslationHypothesis bestHyps = null;
		double phraseItselfScore = 0;
		TranslationHypothesis phraseItself = null;
		for(TranslationHypothesis h : hyps){
			double score = getHypothesisScore(h);
			double length = lengthScore.getNumberOfOptionsForHpy(h);
			if(length == 1){
				phraseItselfScore = score;
				phraseItself = h;
				continue; //it is the phrase p itself
			}
			if(score>bestHypsScore){
				bestHypsScore=score;
				bestHyps = h;
			}
		}
		
		if(bestHyps == null){
			return false; // no way to replicate phrase using smaller phrases
		}
		
		if(bestHypsScore >= phraseItselfScore){
			return true; //if the phrase itself is worse than parts there is no point in keeping it since it won't be used
			//although we should also check if the reorderings orientations are the same 
		}
		
		double relevance = bestHypsScore / phraseItselfScore;
		//System.out.println(phraseItselfScore + "-" + bestHypsScore + "=" + loss);
		return relevance < threshold;
	}

	protected double getHypothesisScore(TranslationHypothesis hyp){
		LinkedList<TranslationOption> options = hyp.getUsedOptions();
		double HypothesisProb = 1;
		double HypothesisRProb = 1;
		double HypothesisLex = 1;
		double HypothesisRLex = 1;
		double HypothesisPenalty = 1;
		
		double HypothesisDistanceDistortion = 0;
		double HypothesisDistortionNext = 1;
		double hypothesisDistortionPrev = 1;
		
		for(int i = 0; i < options.size(); i++){
			TranslationOption option = options.get(i);
			Phrase p = option.getPhrase();
			double prob = ((ProbabilityFeature)p._features.getFeature(PROB_HEADER)).getProbability();
			double rprob = ((ProbabilityFeature)p._features.getFeature(RPROB_HEADER)).getProbability();
			double lex = ((LexicalWeightingFeature)p._features.getFeature(LEXICAL_WEIGHITING_HEADER)).getLexicalWeighting();
			double rlex = ((LexicalWeightingFeature)p._features.getFeature(REVERSE_LEXICAL_WEIGHTING_HEADER)).getLexicalWeighting();
			double penalty = ((PenaltyFeature)p._features.getFeature(PENALTY_HEADER)).getPenalty();
			
			HypothesisProb*=Math.pow(prob,translationWeights[0]);
			HypothesisRProb*=Math.pow(rprob,translationWeights[2]);
			HypothesisLex*=Math.pow(lex,translationWeights[1]);
			HypothesisRLex*=Math.pow(rlex,translationWeights[3]);
			HypothesisPenalty*=Math.pow(penalty,translationWeights[4]);
			
			ReorderingFeature lexReorderingFeatures = ((ReorderingFeature)p._features.getFeature(REORDERING_HEADER));
			TranslationOption prevOpt = null;
			if(i==0){
				prevOpt = null;
			}
			else{
				prevOpt = options.get(i-1);
			}
			int prevOrientation = getOrientation(prevOpt, option);
			hypothesisDistortionPrev*=Math.pow(lexReorderingFeatures.get(0, prevOrientation),reorderingWeights[1 + prevOrientation]);
			
			TranslationOption nextOpt = null;
			if(i==options.size() - 1){
				nextOpt = null;
			}
			else{
				options.get(i+1);
			}
			int nextOrientation = getOrientation(option, nextOpt);
			HypothesisDistortionNext*=Math.pow(lexReorderingFeatures.get(1, nextOrientation), reorderingWeights[1 + 3 + nextOrientation]);
			HypothesisDistanceDistortion-=Math.pow(getDistance(option, nextOpt), reorderingWeights[0]);
		}
		return HypothesisProb
		+ HypothesisLex
		+ HypothesisRProb
		+ HypothesisRLex
		+ HypothesisPenalty
		+ HypothesisDistortionNext
		+ hypothesisDistortionPrev
		+ HypothesisDistanceDistortion;
	}
	
	public int getOrientation(TranslationOption option1, TranslationOption option2){
		if(option1 == null || option2 == null){
			return MSDOrientationCodes.MONO_CODE;
		}
		if(option1.getTargetE() + 1 != option2.getTargetS()){
			return MSDOrientationCodes.OTHER_CODE;
		}
		if(option1.getSourceE() + 1 == option2.getSourceS()){
			return MSDOrientationCodes.MONO_CODE;
		}
		if(option2.getSourceE() + 1 == option1.getSourceS()){
			return MSDOrientationCodes.SWAP_CODE;
		}
		return MSDOrientationCodes.OTHER_CODE;
	}
	
	public int getDistance(TranslationOption option1, TranslationOption option2){
		if(option1 == null || option2 == null){
			return 0;
		}
		return Math.abs(option1.getSourceE() + 1 - option2.getSourceS());
	}
	
	public static LossBasedPrunner parseFromString(String regex, TranslationGraphResource translationGraphResource){
		String[] parts = regex.split(";");
		double threshold = Double.parseDouble(parts[0]);
		String[] translationFeaturesString = parts[1].substring(1, parts[1].length() - 1).split(",");
		String[] reorderingFeaturesString = parts[2].substring(1, parts[2].length() - 1).split(",");
		double[] translationFeatures = new double[translationFeaturesString.length];
		double[] reorderingFeatures = new double[reorderingFeaturesString.length];
		for(int i = 0; i < translationFeatures.length; i++){
			translationFeatures[i] = Double.parseDouble(translationFeaturesString[i]);
		}
		for(int i = 0; i < reorderingFeatures.length; i++){
			reorderingFeatures[i] = Double.parseDouble(reorderingFeaturesString[i]);
		}
		return new LossBasedPrunner(translationGraphResource, translationFeatures, reorderingFeatures, threshold);
	}
}
